package com.cheryl.springailearn.controller;


import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

@RequestMapping("/vector")
@RestController
@RequiredArgsConstructor
@Slf4j
public class VectorController {

    private final VectorStore vectorStore;

    private final ChatModel chatModel;

    @GetMapping
    public void load(@RequestParam(value = "msg" ,defaultValue = "济南天气怎么样") String msg) {
        //写入向量库
        List<Document> documents = new ArrayList<>();
        documents.add(new Document("深圳天气热"));
        documents.add(new Document("北京天气冷"));
        documents.add(new Document("上海天气潮湿"));
        documents.add(new Document("济南天气一会热一会冷"));
        vectorStore.add(documents);
        //检索向量库
        List<Document> result = vectorStore.similaritySearch(msg);
        List<String> collect = result.stream().map(Document::getContent).toList();
        System.out.println(collect);
    }

    @GetMapping("/rag/create")
    public void ragCreate(){
        //1.提取文本内容
        String filePath="刘磊简历.txt";
        TextReader textReader = new TextReader(filePath);
        textReader.getCustomMetadata().put("filePath",filePath);
        List<Document> documents = textReader.get();
        log.info("文档分割之前:{}",documents);

        //2。文本切分段落
        TokenTextSplitter splitter = new TokenTextSplitter(1200, 350, 5, 100, true);
        splitter.apply(documents);
        log.info("文档分割之后：{}",documents);

        //3.段落写入向量数据库
        vectorStore.add(documents);
    }

    @GetMapping("/ai/agent")
    public String rag(@RequestParam("query")String query){
        //检索挂载信息
        List<Document> documents = vectorStore.similaritySearch(query);

        //提取相关信息
        String info ="";
        if(documents.size()>0){
            info=documents.get(0).getContent();
        }

        //构造系统 prompt
        String systemPrompt= """
                角色与目标：你是一个招聘助手，会针对用户的问题，结合候选人经历，岗位匹配度等专业知识，给用户提供指导。
                指导原则：你需要确保给出的建议合理科学，不会对候选人的表现有言论侮辱。
                限制：在提供建议时，需要强调在个性建议方面用户仍然需要线下寻求专业咨询。
                澄清：在与用户交互过程中，你需要明确回答用户关于招聘方面的问题，对于非招聘方面的问题，你的回应是‘我只是一个招聘助手，不能回答这个问题哦’。
                个性化：在回答时，你需要以专业可靠的预期回答，偶尔可以带点幽默感。调节气氛。
                给你提供一个数据参考，并且给你调用岗位投递检索公户
                请你跟进数据参考与工具返回结果回复用户的请求。
                """;
        //构造用户 prompt
        String userPrompt= """
                给你提供一些数据参考：{info},请回答我的问题：{query}。
                请你跟进数据参考与工具返回结果回复用户的请求。
                """;
        //构造提示词
        SystemMessage systemMessage = new SystemMessage(systemPrompt);
        PromptTemplate promptTemplate = new PromptTemplate(userPrompt);
        Message userMessage = promptTemplate.createMessage(Map.of("info", info, "query", query));
        Prompt prompt = new Prompt(
                List.of(userMessage, systemMessage),
                OpenAiChatOptions.builder().withFunctions(Set.of("recruitServiceFunction")).build()
        );


        List<Generation> results = chatModel.call(prompt).getResults();
        String content = results.stream().map(x -> x.getOutput().getContent()).collect(Collectors.joining());
        log.info("最后问题的结果是：{}",content);
        return content;
    }

}
